#!/usr/bin/env python

#########################################################################################
#                                                                                       #
# iQmerge software tool for merging CID and HCD scans for iTRAQ/TMT type experiments    #                                                                            #
# Written by Ashoka D. Polpitiya                                                        #
# for the Translational Genomics Research Institute (TGen), Phoenix, AZ                 #
#                                                                                       #
# Copyright 2011, Translational Genomics Research Institute                             #
# E-mail: ashoka@tgen.org                                                               #
# Website: http://iqmerge.googlecode.com                                                #
# --------------------------------------------------------------------------------------#
#                                                                                       #
# Licensed under the Apache License, Version 2.0 (the "License");                       #
# you may not use this file except in compliance with the License.                      #
# You may obtain a copy of the License at                                               #
#                                                                                       #
#       http://www.apache.org/licenses/LICENSE-2.0                                      #
#                                                                                       #
# Unless required by applicable law or agreed to in writing, software                   #
# distributed under the License is distributed on an "AS IS" BASIS,                     #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.              #
# See the License for the specific language governing permissions and                   #
# limitations under the License.                                                        #
#                                                                                       #
#########################################################################################

from array import array
import sys
from numpy import NaN, Inf, arange, isscalar, asarray

class Merge(object):
    '''
    Class to merge a CID scan with a corresponding HCD scan
    '''

    def __init__(self,sList=None,cidL=None,hcdL=None):
        '''
        Constructor
        '''
        self.scans = sList
        self.cidScans = cidL
        self.hcdScans = hcdL
        
    def mergeCIDHCD(self,sList,dList,cidL,hcdL):
        '''
        Merge matching scans from a list of spectra.
        Matching is done when the precursor m/z are the same. 
        ex: cidL = [[scan#, precursorMZ],...]
        '''
        scans = sList
        metaS = dList
        eTol = 0.001
        
        for iC in cidL:
            for jH in hcdL:
                if abs(jH[1] - iC[1]) < eTol: # Check precursor
                    cidSpec = scans[iC[0]]
                    cidMeta = metaS[iC[0]]
                    hcdSpec = scans[jH[0]]
                    cidSpec, cidMeta = self.mergeTwoScans(cidSpec,hcdSpec,cidMeta)
                    scans[iC[0]] = cidSpec
                    metaS[iC[0]] = cidMeta
        return scans,metaS
    
    def mergeTwoScans(self,cidS,hcdS,cidMeta,mzCutoff=135,intensityTh=0.0005):
        '''
        Merge a CID and a HCD spectra.
        Filter out the low m/z region and the low intensity peaks 
        from HCD spectra
        '''
        cidPara = self.scanParameters(cidS)
        basePeakIntensity = cidPara["basePeakIntensity"]
        
        #Separate the scan to mz, intensity vectors
        mzC = cidS[::2]; yC = cidS[1::2]
        mzH = hcdS[::2]; yH = hcdS[1::2]
        
        #Filter out low mz region from HCD spectra
        idxmzH = [mzH.index(k) for k in mzH if k > mzCutoff]
        mzH = [mzH[i] for i in idxmzH]
        yH = [yH[i] for i in idxmzH]
        
        # Normalize spectra
        yC = [k/max(yC) for k in yC]
        yH = [k/max(yH) for k in yH]
        
        #Filter out low intensities from HCD spectra (need a better way of
        #finding the cutoff, i.e. noise floor
        #idxyH = [yH.index(k) for k in yH if k > intensityTh]
        #mzH = [mzH[i] for i in idxyH]
        #yH = [yH[i] for i in idxyH]
        
        yH, hidx = self.peakdet(yH,intensityTh)
        if len(yH) > 0:
            mzH = [mzH[i] for i in hidx]
            #append both CID and HCD together.
            mzC.extend(mzH)
            yC.extend(yH)
        
        # Sort both based on mz
        idx = range(len(mzC))
        idx.sort(lambda x,y: cmp(mzC[x],mzC[y]))
        mzC = [mzC[i] for i in idx]
        yC = [basePeakIntensity*yC[i] for i in idx]
        
        # Convert to mz-int format
        mergedScan = [None]*(len(mzC) + len(yC))
        mergedScan[::2] = mzC
        mergedScan[1::2] = yC
        
        cidPara = self.scanParameters(mergedScan)
        cidMeta["basePeakIntensity"] = cidPara["basePeakIntensity"]
        cidMeta['highMz'] = cidPara['highMz']
        cidMeta['lowMz'] = cidPara['lowMz']
        cidMeta['basePeakMz'] = cidPara['basePeakMz']
        cidMeta['peaksCount'] = cidPara['peaksCount']
        
        return array('f',mergedScan),cidMeta
        
                        
    def scanParameters(self,s):
        '''
        Retrieve the Scan specific parameters:
        highMz, lowMz, basePeakMZ, basePeakIntensity, peaksCount
        '''
        totIonCurrent = 0
        maxmz = None
        minmz = None
        basePeakMz = None
        basePeakIntensity = None
        peaksCount = 0
        for i in xrange(0,len(s),2):
            x = s[i]; y = s[i+1]
            if minmz is None or minmz > x:
                minmz = x
            if maxmz is None or maxmz < x:
                maxmz = x
            totIonCurrent += y
            if basePeakIntensity is None or basePeakIntensity < y:
                basePeakIntensity = y
                basePeakMz = x
            peaksCount += 1
            
        out = dict({'highMz':maxmz, 'lowMz':minmz, 'basePeakMz':basePeakMz, \
                    'basePeakIntensity':basePeakIntensity, 'peaksCount':peaksCount})
        return out
        

    def peakdet(self, v, delta):
        """
        Based on the MATLAB script at http://billauer.co.il/peakdet.html
        
        Currently returns two lists of tuples, but maybe arrays would be better
        
        function [maxtab, idx]=peakdet(v, delta)
        %PEAKDET Detect peaks in a vector
        %        [MAXTAB, IDX] = PEAKDET(V, DELTA) finds the local
        %        maxima and their indices ("peaks") in the vector V.
        %      
        %        A point is considered a maximum peak if it has the maximal
        %        value, and was preceded (to the left) by a value lower by
        %        DELTA.
        
        % Eli Billauer, 3.4.05 (Explicitly not copyrighted).
        % This function is released to the public domain; Any use is allowed.
        
        """
        maxtab = []
        maxIdx = []
           
        
        x = arange(len(v))
        
        v = asarray(v)
                
        if not isscalar(delta):
            sys.exit('Input argument delta must be a scalar')
        
        if delta <= 0:
            sys.exit('Input argument delta must be positive')
        
        mn, mx = Inf, -Inf
        mxpos = NaN
        
        lookformax = True
        
        for i in arange(len(v)):
            this = v[i]
            if this > mx:
                mx = this
                mxpos = x[i]
            if this < mn:
                mn = this
            
            if lookformax:
                if this < mx-delta:
                    maxtab.append(mx)
                    maxIdx.append(mxpos)
                    mn = this
                    lookformax = False
            else:
                if this > mn+delta:
                    mx = this
                    mxpos = x[i]
                    lookformax = True
                
        return maxtab, maxIdx
                